use std::collections::HashMap;
use std::fmt::{Debug, Formatter};
use std::ops::Range;
use std::sync::Arc;
use std::time;

use datafusion::arrow::array::{ArrayBuilder, ArrayRef, Float64Builder, GenericStringBuilder, Int64Builder, RecordBatch, UInt32Builder, UInt64Builder};
use datafusion::arrow::datatypes::{DataType, FieldRef, Fields, SchemaRef};
use log::debug;
use parking_lot::RwLock;

use common::{ArrayBuilderExt, Entry, Error, FieldValue, LabelType, LResult, SchemaMut};
use engine::TableScanOption;

pub type MemTableRef = Arc<RwLock<MemTable>>;

pub const DEFAULT_MEM_TABLE_CHUNK_SIZE: usize = 4096;

pub struct MemTable {
    measurement: String,
    // ts, tags..., fields...
    ts_seconds: UInt64Builder,
    ts_nanos: UInt32Builder,
    tags: HashMap<String, Box<dyn ArrayBuilder>>,
    fields: HashMap<String, Box<dyn ArrayBuilder>>,
    pub(crate) schema: SchemaMut,
}

unsafe impl Sync for MemTable {}

impl MemTable {
    pub fn create(measurement: String) -> Self {
        Self {
            measurement,
            ts_seconds: Default::default(),
            ts_nanos: Default::default(),
            tags: Default::default(),
            fields: Default::default(),
            schema: SchemaMut::create(),
        }
    }

    #[inline]
    pub fn measurement(&self) -> &String {
        &self.measurement
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.ts_seconds.len()
    }

    pub fn push(&mut self, entry: Entry) -> LResult<()> {
        self.ensure_schema(&entry)?;
        let Entry { ts, tags, fields, .. } = entry;
        let d = ts.duration_since(time::UNIX_EPOCH).unwrap();
        self.ts_seconds.append_value(d.as_secs());
        self.ts_nanos.append_value(d.subsec_nanos());

        let len = self.ts_seconds.len();
        for (k, v) in tags {
            let key = unsafe { k.get_string_ref() }.to_string();
            let a = self.tags.entry(key).or_insert_with(|| {
                new_tag_value_array(len - 1)
            });
            let b: &mut UInt64Builder = unsafe { a.as_any_mut().downcast_mut_unchecked() };
            b.append_value(v.get_value());
        }
        // 补齐未填充的字段
        for (_, tag) in &mut self.tags {
            if tag.len() < len {
                unsafe {
                    tag.as_any_mut().downcast_mut_unchecked::<UInt64Builder>().append_null();
                }
            }
        }
        // 处理写入的 fields
        for (k, v) in fields {
            let key = unsafe { k.get_string_ref() }.to_string();
            let a = self.fields.entry(key).or_insert_with(|| {
                new_field_value_array(len - 1, &v)
            });
            append_to_field_value(a, v);
        }
        // 补齐未填充的字段
        for (key, field) in &mut self.fields {
            if field.len() < len {
                let arrow_field = self.schema.get_field(&*key).unwrap();
                append_null_to_field_value(field, arrow_field.data_type());
            }
        }
        Ok(())
    }

    fn ensure_schema(&mut self, entry: &Entry) -> LResult<()> {
        for k in entry.tags.keys() {
            let tag = unsafe { &*k.get_string_ref() };
            self.schema.ensure_tag(tag);
        }
        for (k, v) in &entry.fields {
            let field = unsafe { &*k.get_string_ref() };
            self.schema.ensure_field(field, v.data_type());
        }
        Ok(())
    }

    pub fn read(&self, range: Range<usize>, schema: &SchemaRef, option: &TableScanOption) -> LResult<RecordBatch> {
        let mut array: Vec<ArrayRef> = vec![];
        let projection_fields = self.iter_with_projection(schema, option);
        for f in &*projection_fields {
            debug!(target: "engine:mem", "read column: {}", f.name());
            if f.name().eq("ts_s") {
                array.push(self.ts_seconds.slice(&DataType::UInt64, range.clone()));
            } else if f.name().eq("ts_n") {
                array.push(self.ts_nanos.slice(&DataType::UInt32, range.clone()));
            } else {
                let op = match SchemaMut::dismiss_name(f.name())? {
                    (name, LabelType::Tag) => {
                        self.tags.get(name)
                    }
                    (name, LabelType::Field) => {
                        self.fields.get(name)
                    }
                    _ => unreachable!()
                };
                if op.is_none() {
                    return Err(Error::StorageErr(format!("MemTable illegal label: {}", f.name())));
                }
                let builder = op.unwrap();
                array.push(builder.slice(f.data_type(), range.clone()));
            }
        }
        // record_batch schema
        let record_schema = if let Some(projection) = &option.projection {
            Arc::new(schema.project(projection).unwrap())
        } else {
            schema.clone()
        };
        RecordBatch::try_new(record_schema, array).map_err(|e| common::Error::StorageErr(
            format!("MemTable read RecordBatch err: {}", e)
        ))
    }

    fn iter_with_projection(&self, schema: &SchemaRef, option: &TableScanOption) -> Fields {
        if let Some(projection) = &option.projection {
            let projection_fields: Vec<FieldRef> = schema.fields().iter().enumerate().filter(|(i, _)| {
                projection.contains(i)
            }).map(|(_i, f)| f.clone()).collect();
            projection_fields.into()
        } else {
            schema.fields().clone()
        }
    }
}

// fn should_read_ts(schema: &SchemaRef, option: &TableScanOption) -> bool {
//     if let Some(projection) = &option.projection {
//         // ts always at index 0
//         projection[0] == 0
//     } else {
//         true
//     }
// }

fn new_tag_value_array(len: usize) -> Box<dyn ArrayBuilder> {
    let mut b = UInt64Builder::with_capacity(DEFAULT_MEM_TABLE_CHUNK_SIZE);
    b.append_nulls(len);
    Box::new(b)
}


// TODO, make it into a macro
fn new_field_value_array(len: usize, v: &FieldValue) -> Box<dyn ArrayBuilder> {
    match v {
        FieldValue::I64(_) => {
            let mut b = Int64Builder::with_capacity(DEFAULT_MEM_TABLE_CHUNK_SIZE);
            b.append_nulls(len);
            Box::new(b)
        }
        FieldValue::F64(_) => {
            let mut b = Float64Builder::with_capacity(DEFAULT_MEM_TABLE_CHUNK_SIZE);
            b.append_nulls(len);
            Box::new(b)
        }
        FieldValue::Timestamp(_) => {
            let mut b = UInt64Builder::with_capacity(DEFAULT_MEM_TABLE_CHUNK_SIZE);
            b.append_nulls(len);
            Box::new(b)
        }
        FieldValue::String(_) => {
            let mut b = GenericStringBuilder::<i32>::with_capacity(DEFAULT_MEM_TABLE_CHUNK_SIZE, DEFAULT_MEM_TABLE_CHUNK_SIZE);
            (0..len).for_each(|_| {
                b.append_null();
            });
            Box::new(b)
        }
    }
}

fn append_to_field_value(a: &mut Box<dyn ArrayBuilder>, fv: FieldValue) {
    unsafe {
        match fv {
            FieldValue::I64(v) => {
                a.as_any_mut().downcast_mut_unchecked::<Int64Builder>().append_value(v);
            }
            FieldValue::F64(v) => {
                a.as_any_mut().downcast_mut_unchecked::<Float64Builder>().append_value(v);
            }
            FieldValue::Timestamp(v) => {
                a.as_any_mut().downcast_mut_unchecked::<UInt64Builder>().append_value(v);
            }
            FieldValue::String(v) => {
                a.as_any_mut().downcast_mut_unchecked::<GenericStringBuilder<i32>>().append_value(v);
            }
        }
    }
}

fn append_null_to_field_value(a: &mut Box<dyn ArrayBuilder>, dt: &DataType) {
    unsafe {
        match dt {
            &DataType::Int64 => {
                a.as_any_mut().downcast_mut_unchecked::<Int64Builder>().append_null();
            }
            &DataType::Float64 => {
                a.as_any_mut().downcast_mut_unchecked::<Float64Builder>().append_null();
            }
            &DataType::Timestamp(_, _) => {
                a.as_any_mut().downcast_mut_unchecked::<UInt64Builder>().append_null();
            }
            &DataType::Utf8 => {
                a.as_any_mut().downcast_mut_unchecked::<GenericStringBuilder<i32>>().append_null();
            }
            other => {
                unimplemented!("append_null_to_field_value DataType: {}", other)
            }
        }
    }
}

impl Debug for MemTable {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}({})", self.measurement, self.len())
    }
}

#[cfg(test)]
mod test {
    use std::sync::Arc;

    use common::{Entry, FieldValue, LResult};

    use crate::mem_table::MemTable;

    #[test]
    fn test() -> LResult<()> {
        let mut mem_table = MemTable::create("a".into());
        {
            let e = Entry::new_builder("a")
                .add_tag("tag_a".try_into()?, "a".try_into()?)
                .add_field("field_a".try_into()?, FieldValue::I64(12))
                .add_field("field_c".try_into()?, FieldValue::String("aaa".into()))
                .build()?;
            mem_table.push(e)?;
        }
        {
            let e = Entry::new_builder("a")
                .add_tag("tag_b".try_into()?, "a".try_into()?)
                .add_field("field_b".try_into()?, FieldValue::I64(12))
                .build()?;
            mem_table.push(e)?;
        }
        assert_eq!(2, mem_table.len());
        // assert schema
        assert_eq!(7, mem_table.schema.fields.len());
        assert!(mem_table.schema.contains_real_field("ts_s"));
        assert!(mem_table.schema.contains_real_field("ts_n"));
        assert!(mem_table.schema.contains_real_field("$t_tag_a"));
        assert!(mem_table.schema.contains_real_field("$t_tag_b"));
        assert!(mem_table.schema.contains_real_field("$f_field_a"));
        assert!(mem_table.schema.contains_real_field("$f_field_b"));
        assert!(mem_table.schema.contains_real_field("$f_field_c"));

        let schema = Arc::new(mem_table.schema.as_schema());
        let records = mem_table.read(0..2, &schema, Default::default())?;
        assert_eq!(2, records.num_rows());
        assert_eq!(7, records.num_columns());

        Ok(())
    }
}